#!/usr/bin/env python3
# -*- coding: utf-8 -*-

import sys
import rospy
import numpy as np
import tf
from sensor_msgs.msg import JointState
from geometry_msgs.msg import Pose, Point
from std_msgs.msg import Float64
import moveit_msgs.msg
import geometry_msgs.msg
import tf.transformations
import moveit_commander
from moveit_commander import MoveGroupCommander
from copy import deepcopy
import time
from moveit_commander import RobotTrajectory
from trajectory_msgs.msg import JointTrajectoryPoint

sign = False
joint_state_data = None
result2_data = None
tf_viewer = None
allow_visual_control = False
class PIDController:
    def __init__(self, kp, ki, kd):
        self.kp = kp
        self.ki = ki
        self.kd = kd
        self.prev_error = 0
        self.integral = 0
        self.last_time = time.time()

    def compute(self, setpoint, measured_value):
        current_time = time.time()
        dt = current_time - self.last_time
        if dt <= 0.0:
            return 0

        error = setpoint - measured_value
        self.integral += error * dt
        derivative = (error - self.prev_error) / dt

        output = self.kp * error + self.ki * self.integral + self.kd * derivative

        self.prev_error = error
        self.last_time = current_time

        return output
    
def scale_trajectory_speed(traj, scale):
    # 创建新的轨迹对象并初始化
    new_traj = RobotTrajectory()
    new_traj.joint_trajectory = traj.joint_trajectory

    # 获取关节和轨迹点数量
    n_joints = len(traj.joint_trajectory.joint_names)
    n_points = len(traj.joint_trajectory.points)
    points = list(traj.joint_trajectory.points)
    
    # 调整时间、速度和加速度
    for i in range(n_points):
        point = JointTrajectoryPoint()
        point.positions = traj.joint_trajectory.points[i].positions
        point.time_from_start = traj.joint_trajectory.points[i].time_from_start / scale
        point.velocities = list(traj.joint_trajectory.points[i].velocities)
        point.accelerations = list(traj.joint_trajectory.points[i].accelerations)
        
        for j in range(n_joints):
            point.velocities[j] *= scale
            point.accelerations[j] *= scale * scale

        points[i] = point

    new_traj.joint_trajectory.points = points
    return new_traj


def joint_states_callback(data):
    global joint_state_data
    joint_state_data = data

def result2_callback(data):
    global result2_data
    result2_data = data

def process_data(event):
    global joint_state_data, result2_data, tf_viewer, allow_visual_control, sign
    if sign:
        raise KeyboardInterrupt

    if True:
            print("进入计算")
            joint_positions = [joint_state_data.position[i] if i < len(joint_state_data.position) else 0 for i in range(5)]
            joint_positions_degrees = [np.degrees(theta) for theta in joint_positions]
            joint_positions_degrees_rounded = [round(deg) for deg in joint_positions_degrees]
            print("xyz:",result2_data.x,result2_data.y,result2_data.z)

            # 定义theta变量
            theta1, theta2, theta3, theta4, theta5 = joint_positions
            theta6 = 0

            cos = np.cos
            sin = np.sin

            TS = np.array([
            [  cos(theta5)*sin(theta1) - sin(theta5)*(cos(theta4)*(cos(theta1)*sin(theta2)*sin(theta3) - cos(theta1)*cos(theta2)*cos(theta3)) + sin(theta4)*(cos(theta1)*cos(theta2)*sin(theta3) + cos(theta1)*cos(theta3)*sin(theta2))), sin(theta4)*(cos(theta1)*sin(theta2)*sin(theta3) - cos(theta1)*cos(theta2)*cos(theta3)) - cos(theta4)*(cos(theta1)*cos(theta2)*sin(theta3) + cos(theta1)*cos(theta3)*sin(theta2)), - sin(theta1)*sin(theta5) - cos(theta5)*(cos(theta4)*(cos(theta1)*sin(theta2)*sin(theta3) - cos(theta1)*cos(theta2)*cos(theta3)) + sin(theta4)*(cos(theta1)*cos(theta2)*sin(theta3) + cos(theta1)*cos(theta3)*sin(theta2))), 150*cos(theta1)*sin(theta2) - (8693*sin(theta1)*sin(theta5))/100 + (77*cos(theta4)*(cos(theta1)*cos(theta2)*sin(theta3) + cos(theta1)*cos(theta3)*sin(theta2)))/2 - (141*cos(theta4)*(cos(theta1)*sin(theta2)*sin(theta3) - cos(theta1)*cos(theta2)*cos(theta3)))/2 - (141*sin(theta4)*(cos(theta1)*cos(theta2)*sin(theta3) + cos(theta1)*cos(theta3)*sin(theta2)))/2 - (77*sin(theta4)*(cos(theta1)*sin(theta2)*sin(theta3) - cos(theta1)*cos(theta2)*cos(theta3)))/2 - (8693*cos(theta5)*(cos(theta4)*(cos(theta1)*sin(theta2)*sin(theta3) - cos(theta1)*cos(theta2)*cos(theta3)) + sin(theta4)*(cos(theta1)*cos(theta2)*sin(theta3) + cos(theta1)*cos(theta3)*sin(theta2))))/100 + 150*cos(theta1)*cos(theta2)*sin(theta3) + 150*cos(theta1)*cos(theta3)*sin(theta2) - 483/4],
            [- cos(theta1)*cos(theta5) - sin(theta5)*(cos(theta4)*(sin(theta1)*sin(theta2)*sin(theta3) - cos(theta2)*cos(theta3)*sin(theta1)) + sin(theta4)*(cos(theta2)*sin(theta1)*sin(theta3) + cos(theta3)*sin(theta1)*sin(theta2))), sin(theta4)*(sin(theta1)*sin(theta2)*sin(theta3) - cos(theta2)*cos(theta3)*sin(theta1)) - cos(theta4)*(cos(theta2)*sin(theta1)*sin(theta3) + cos(theta3)*sin(theta1)*sin(theta2)),   cos(theta1)*sin(theta5) - cos(theta5)*(cos(theta4)*(sin(theta1)*sin(theta2)*sin(theta3) - cos(theta2)*cos(theta3)*sin(theta1)) + sin(theta4)*(cos(theta2)*sin(theta1)*sin(theta3) + cos(theta3)*sin(theta1)*sin(theta2))),         (8693*cos(theta1)*sin(theta5))/100 + 150*sin(theta1)*sin(theta2) + (77*cos(theta4)*(cos(theta2)*sin(theta1)*sin(theta3) + cos(theta3)*sin(theta1)*sin(theta2)))/2 - (141*cos(theta4)*(sin(theta1)*sin(theta2)*sin(theta3) - cos(theta2)*cos(theta3)*sin(theta1)))/2 - (141*sin(theta4)*(cos(theta2)*sin(theta1)*sin(theta3) + cos(theta3)*sin(theta1)*sin(theta2)))/2 - (77*sin(theta4)*(sin(theta1)*sin(theta2)*sin(theta3) - cos(theta2)*cos(theta3)*sin(theta1)))/2 - (8693*cos(theta5)*(cos(theta4)*(sin(theta1)*sin(theta2)*sin(theta3) - cos(theta2)*cos(theta3)*sin(theta1)) + sin(theta4)*(cos(theta2)*sin(theta1)*sin(theta3) + cos(theta3)*sin(theta1)*sin(theta2))))/100 + 150*cos(theta2)*sin(theta1)*sin(theta3) + 150*cos(theta3)*sin(theta1)*sin(theta2)],
            [                                                                           -sin(theta5)*(cos(theta4)*(cos(theta2)*sin(theta3) + cos(theta3)*sin(theta2)) + sin(theta4)*(cos(theta2)*cos(theta3) - sin(theta2)*sin(theta3))),                                                 sin(theta4)*(cos(theta2)*sin(theta3) + cos(theta3)*sin(theta2)) - cos(theta4)*(cos(theta2)*cos(theta3) - sin(theta2)*sin(theta3)),                                                                            -cos(theta5)*(cos(theta4)*(cos(theta2)*sin(theta3) + cos(theta3)*sin(theta2)) + sin(theta4)*(cos(theta2)*cos(theta3) - sin(theta2)*sin(theta3))),                                                                                                                                                                                                                          150*cos(theta2) + 150*cos(theta2)*cos(theta3) - 150*sin(theta2)*sin(theta3) - (8693*cos(theta5)*(cos(theta4)*(cos(theta2)*sin(theta3) + cos(theta3)*sin(theta2)) + sin(theta4)*(cos(theta2)*cos(theta3) - sin(theta2)*sin(theta3))))/100 - (141*cos(theta4)*(cos(theta2)*sin(theta3) + cos(theta3)*sin(theta2)))/2 + (77*cos(theta4)*(cos(theta2)*cos(theta3) - sin(theta2)*sin(theta3)))/2 - (77*sin(theta4)*(cos(theta2)*sin(theta3) + cos(theta3)*sin(theta2)))/2 - (141*sin(theta4)*(cos(theta2)*cos(theta3) - sin(theta2)*sin(theta3)))/2 + 193/2],
            [                                                                                                                                                                                                                          0,                                                                                                                                                                                 0,                                                                                                                                                                                                                           0,                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                               1]
            ])   
            # rospy.loginfo('TS:\n %s', TS)

            global M
            M = np.eye(4) # 定义M矩阵
            M[0, 3] = -result2_data.x

            M[1, 3] =result2_data.y -16+5

            M[2, 3] =  result2_data.z -(161-86.93) +50 +5

            Tc12 = np.array([[0, -1, 0, 0], [0, 0, -1, 0], [1, 0, 0, 0], [0, 0, 0, 1]])
            Tce = np.array([[1, 0, 0, 88.876], [0, 1, 0, 0], [0, 0, 1, -59.6856], [0, 0, 0, 1]])
            result_matrix = TS @ M @ Tc12 
            result_matrix[:3, 3] /= 1000  # 调整尺度
            # 对矩阵中的值进行四舍五入
            rounded_result_matrix = np.round(result_matrix, decimals=3)  # 四舍五入到小数点后 6 位
            # 打印四舍五入后的矩阵
            # rospy.loginfo('Rounded Result Matrix:\n %s', rounded_result_matrix)
            position = result_matrix[:3, 3] # 提取位置和旋转矩阵
            rotation_matrix = result_matrix[:3, :3]
            # 打印位置和旋转矩阵
            print("Position:", position)
            # print("Rotation Matrix:\n", rotation_matrix)
            quaternion = tf.transformations.quaternion_from_matrix(result_matrix) # 将旋转矩阵转换为四元数
            # # 打印位置和四元数
            print("Quaternion:\n", quaternion)
            
            # success = move_arm_to_position(position, quaternion)
            target_pose = arm.get_current_pose().pose
            target_pose.position.x = position[0]
            target_pose.position.y = position[1]
            target_pose.position.z = position[2]
            target_pose.orientation.x = quaternion[0]
            target_pose.orientation.y = quaternion[1]
            target_pose.orientation.z = quaternion[2]
            target_pose.orientation.w = quaternion[3]
            
            # 设置目标姿态
            arm.set_pose_target(target_pose)
            waypoints = [target_pose]

            (plan, fraction) = arm.compute_cartesian_path(waypoints,0.01,False)
            new_plan = scale_trajectory_speed(plan, 0.5)
            success = arm.execute(new_plan, wait=True)

            if success:
                sign=True
                # rospy.signal_shutdown("Shutting down node after successful motion plan")
                print("退出视觉")

def move_arm_to_position(position, quaternion):
    global arm
    global M

    # 获取当前机械臂的位置
    current_pose = arm.get_current_pose().pose

    # 计算当前位置与目标位置的距离
    current_position = np.array([current_pose.position.x, current_pose.position.y, current_pose.position.z])
    target_position = np.array(position)
    distance = np.linalg.norm(current_position - target_position)

    # 定义阈值（例如，0.01米）
    threshold = 0.07
    if distance < threshold:
        print("当前位置与目标位置相差不远，不做规划。")
        return  False # 取消规划


    target_pose = Pose()  # 设置目标位姿
    target_pose.position.x = target_position[0]
    target_pose.position.y = target_position[1]
    target_pose.position.z = target_position[2]
    target_pose.orientation.x = quaternion[0]
    target_pose.orientation.y = quaternion[1]
    target_pose.orientation.z = quaternion[2]
    target_pose.orientation.w = quaternion[3]

    # 控制机械臂运动到目标位姿
    arm.set_pose_target(target_pose)
    success = arm.go()
    rospy.sleep(1)

    if success:
        if M[0, 3] == 0 and M[1, 3] == 0 and M[2, 3] == 0:
            print("保持不动\n")
        else:
            print("规划成功\n")
        return True
    else:
        print("未规划")
        return False
allow_visual_control = True  # 允许视觉

def stop_visual_control2():
    global allow_visual_control
    allow_visual_control = False
    sign = True

def run_visual_control2():
    global allow_visual_control, sign
    allow_visual_control = True
    sign = False

def visual_once():
    global arm
    print("视觉开始")
    # rospy.init_node('multi_topic_listener', anonymous=True)  # 保留这个初始化节点的调用

    moveit_commander.roscpp_initialize(sys.argv) 
    
    arm = MoveGroupCommander('manipulator') 

    arm.allow_replanning(True) 

    arm.set_pose_reference_frame('base_link') 

    arm.set_goal_position_tolerance(0.0001) 
    arm.set_goal_orientation_tolerance(0.0001)

    arm.set_max_acceleration_scaling_factor(0.1)
    arm.set_max_velocity_scaling_factor(0.1)

    # 控制机械臂先回到初始化位置
    arm.set_named_target('middle')
    arm.go()
    rospy.sleep(1)

    global tf_viewer
    tf_viewer = rospy.Publisher('result_data', Float64, queue_size=10)
    rospy.Subscriber("/joint_states", JointState, joint_states_callback)
    rospy.Subscriber("/result3_topic", Point, result2_callback)

    # ## 模拟接收到 result2_topic 消息
    # global result2_data
    # result2_data = Point(0, 0, 0)

    rospy.Timer(rospy.Duration(1), process_data)
    # rospy.spin()


if __name__ == '__main__':
    try:
        moveit_commander.roscpp_initialize(sys.argv)
        visual_once()
    except rospy.ROSInterruptException:
        pass